"""
负责存放阿里的通义模型的调用逻辑
"""
import base64
import concurrent
import os
import tempfile
from pathlib import Path
from typing import TypedDict

import dashscope
from dashscope import MultiModalConversation, Generation
from qwen3_asr_toolkit import qwen3asr
from qwen3_asr_toolkit.audio_tools import load_audio, WAV_SAMPLE_RATE, process_vad, save_audio_file
from silero_vad import load_silero_vad


class ModelScopeError(Exception):
    pass

class _VideoFrameSlice(TypedDict):
    """
    视频帧截图文件路径列表
    """
    frames: list[Path]
    """
    视频帧分片起始时间点
    """
    start_seconds: float
    """
    视频帧分片结束时间点
    """
    end_seconds: float

class ModelScopeService:

    def __init__(self, api_key = None):
        if api_key is None:
            env_key = os.getenv("DASHSCOPE_API_KEY")
            if env_key is None:
                raise ValueError("api_key is None")

            self.api_key = env_key
        else:
            self.api_key = api_key


    def chat(self, prompt: str, temperature: float = 0.8, use_json_format_response: bool = False):
        """
        调用模型服务，进行对话
        :param prompt: 对话内容
        :param temperature: 采样温度
        :param use_json_format_response: 是否使用json格式的返回值
        :return:
        """
        if not prompt:
            raise ValueError("prompt is required")

        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "text": prompt
                    }
                ]
            }
        ]

        response = Generation.call(
            api_key = self.api_key,
            model = "qwen-flash",
            messages = messages,
            stream = False,
            result_format = "message",
            temperature = temperature,
            response_format = {"type" : "json_object"} if use_json_format_response else None,
        )

        if response.status_code != 200:
            raise ModelScopeError(f"调用modelscope api时发生错误,code={response.code}, msg={response.message}")

        return response.output.choices[0].message.content


    def transcript_audio(self, audio_file_path: Path):
        """
        调用模型服务，将语音文件转写为文本
        :param audio_file_path: 文件路径
        :return:
        """
        if not audio_file_path.is_file():
            raise ValueError("audio_file_path is not a file")

        # transcode_target_path = audio_file_path.parent / str(uuid.uuid4())
        # ffmpeg_tools.audio_transcode(audio_file_path, transcode_target_path,
        #                              sample_rate = 16000, channels = 1, sample_format = "s16")

        # 下面代码借鉴自qwen3 asr toolkit的call_api.py，因为qwen3-asr-toolkit只提供了命令行调用说明，代码里怎么调用只能模仿它命令行模式的代码了
        # toolkit里的load_audiof函数自带转码处理了
        wav = load_audio(audio_file_path)

        # 根据音频的bytes数和每秒采样率进行计算时长，超过3分钟就调用vad模型来对饮品进行切分
        if len(wav) / WAV_SAMPLE_RATE >= 180:
            worker_vad_model = load_silero_vad(onnx=True)
            wav_list = process_vad(wav, worker_vad_model, segment_threshold_s=120, max_segment_threshold_s=180)
        else:
            wav_list = [(0, len(wav), wav)]

        # 把切分后的文件找个临时路径存起来
        wav_name = os.path.basename(audio_file_path)
        wav_dir_name = os.path.splitext(wav_name)[0]
        save_dir = os.path.join(tempfile.gettempdir(), wav_dir_name)

        wav_path_list = []
        for idx, (_, _, wav_data) in enumerate(wav_list):
            wav_path = os.path.join(save_dir, f"{wav_name}_{idx}.wav")
            save_audio_file(wav_data, wav_path)
            wav_path_list.append(wav_path)

        # Multithread call qwen3-asr-flash api
        results = []
        with concurrent.futures.ThreadPoolExecutor(max_workers = 4) as executor:
            future_dict = {
                executor.submit(self._call_audio_transcript_service, wav_path): idx
                for idx, wav_path in enumerate(wav_path_list)
            }

            for future in concurrent.futures.as_completed(future_dict):
                idx = future_dict[future]

                try :
                    results.append((idx, future.result()))
                except Exception as e:
                    print(f"Error processing wav file {wav_path}: {e}")
                    continue

        # Sort and splice in the original order
        results.sort(key=lambda x: x[0])
        full_text = " ".join(text for _, text in results)

        os.system(f"rm -rf {save_dir}")

        return full_text


    def video_understanding(self, video_frame_list: list[Path], fps: float):
        """
        调用模型服务，执行视频内容的理解，并转写为文本
        :param video_frame_list: 视频帧截图文件路径列表
        :param fps: 生成视频帧图片文件时使用的fps
        :return: 第一个返回值为最终的分析结果，第二个返回值为分析过程中每个批次的分析结果
        """

        if not video_frame_list or len(video_frame_list) == 0:
            raise ValueError("video_file_path is required")
        if not fps:
            raise ValueError("fps is required")

        video_frame_parts = self._slice_video_frame(video_frame_list, fps)

        batch_results = []
        for i in range(len(video_frame_parts)):
            if i > 0:
                batch_results.append(self._video_understanding_batch(video_frame_parts[i], fps, batch_results[i - 1]))
            else:
                batch_results.append(self._video_understanding_batch(video_frame_parts[i], fps, None))

        return batch_results[len(batch_results) - 1], batch_results


    def translate_bbox_to_image_coord(self, bbox: tuple[int, int, int, int], image_coord: tuple[int, int]) -> tuple[int, int, int, int]:
        """
        将LLM返回的坐标数据，转换为以图像的分辨率为基地的坐标值
        :param bbox: llm返回的坐标数值，为固定4个元素的tuple
        :param image_coord: 图片的横向长度与纵向长度组成的tuple
        :return: 转换后的坐标数值，为固定4个元素的tuple
        """

        # qwen3-vl会将所有图片的坐标归一化到[0,999]，所以这里转化的基数取1000
        horizontal_factor = image_coord[0] / 1000
        vertical_factor = image_coord[1] / 1000

        return (int(bbox[0] * horizontal_factor),
                  int(bbox[1] * vertical_factor),
                  int(bbox[2] * horizontal_factor),
                  int(bbox[3] * vertical_factor))


    def _call_audio_transcript_service(self, audio_file_path: Path):
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "audio": f"file://{str(audio_file_path)}"
                    }
                ]
            }
        ]

        response = dashscope.MultiModalConversation.call(
            api_key = self.api_key,
            model = "qwen3-asr-flash",
            messages = messages,
            result_format = "message",
            asr_options = {
                "enable_itn": True
            }
        )

        if response.status_code != 200:
            raise ModelScopeError(f"调用modelscope api时发生错误,code={response.code}, msg={response.message}")
        else:
            return response.output.choices[0].message.content[0]["text"]


    def _slice_video_frame(self, video_frame_list: list[Path], fps: float, max_slice_size: int = 250) -> list[_VideoFrameSlice]:
        """
        对视频抽帧图片文件进行分片
        :param video_frame_list:
        :param fps:
        :return:
        """

        ret = []
        for i in range(0, len(video_frame_list), max_slice_size):
            frame_size = min(max_slice_size, len(video_frame_list) - i)
            ret.append(_VideoFrameSlice(frames=video_frame_list[i : i + frame_size],
                                        start_seconds=i / fps,
                                        end_seconds=(i + frame_size) / fps))

        return ret;

    def _video_understanding_batch(self, frame_slice: _VideoFrameSlice, fps: float, prev_context: str | None):
        """
        调用模型服务，执行视频内容的理解，并转写为文本
        :param frame_slice: 视频帧截图文件路径列表
        :param fps: 帧率
        :param prev_context: 要添加作为上下文的信息
        :return:
        """

        # https://bailian.console.aliyun.com/?spm=5176.29619931.J_PvCec88exbQTi-U433Fxg.5.74cd405fkjktLa&tab=doc#/doc/?type=model&url=2845871

        # frame_img_path_list = [f"file://{str(x.absolute())}" for x in video_frame_list]
        #
        # messages = [
        #     {
        #         'role': 'user',
        #         'content': [
        #             {
        #                 'video': frame_img_path_list,
        #                 'fps': fps
        #             },
        #             {
        #                 'text': self._build_video_understanding_prompt(prev_context)
        #             }
        #         ]
        #     }
        # ]
        #
        encoded_img_list = [f"data:image/jpeg;base64,{self._encode_image(x)}" for x in frame_slice["frames"]]
        submit_msg = [
            {
                'role': 'user',
                'content': [
                    {
                        "type": "video",
                        "video": [
                            *encoded_img_list
                        ],
                        "fps": fps
                    },
                    {
                        "type": "text",
                        "text": self._build_video_understanding_prompt(prev_context, current_frame_slice = frame_slice)
                    }
                ]
            }
        ]


        response = MultiModalConversation.call(
            api_key=self.api_key,
            model='qwen3-vl-plus', # qwen3-vl-flash在视频元素的提取上做的很差，时间点基本上不准确，且提取的元素质量也不高
            messages=submit_msg,
            response_format={'type': 'json_object'}
        )

        if response.status_code != 200:
            raise ModelScopeError(f"调用modelscope api时发生错误,code={response.code}, msg={response.message}")
        else:
            return response.output.choices[0].message.content[0]["text"]

    def _encode_image(self, image_path: Path):
        """
        对图片文件进行base64编码
        :return:
        """
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode("utf-8")

    def _build_video_understanding_prompt(self, prev_context: str | None, current_frame_slice: _VideoFrameSlice):
        """
        构建视频理解任务的prompt
        :param prev_context: 上一个视频切片分析的结果
        :param current_frame_slice: 当前视频切片
        :return:
        """
        # if prev_context:
        #     return f"""
        #     该视频为一个长视频的切片，请结合之前的视频切片的分析结果和当前视频切片的信息，将视频的内容进行理解，转写为文本，
        #     并且需要保留较为详细的信息，以便下一个视频切片在总结视频内容时有足够的信息。
        #     同时如果某个时间点的视频画面中存在包含了重要信息的元素的话，将其标注出来，为每个物体生成边界框，返回对应时间点以及[x_min, y_min, x_max, y_max]格式的坐标，而不只是用文本对其内容进行声明。
        #     此前的视频切片的分析结果为为：
        #     {prev_context}
        #     """

        prompt = f"""
该视频为一个长视频的切片，时间范围是第{current_frame_slice['start_seconds']}秒到第{current_frame_slice['end_seconds']}秒。

# 任务要求
1. **视频内容总结**：请对该视频切片进行详细的内容总结，保留关键信息，确保下一个视频切片在总结时有足够的上下文信息。

2. **精确元素定位**：对于总结中提到的重要视觉元素，需要进行精确的空间定位（grounding

# 元素定位的精确要求
## 边界框准确性
- **完整覆盖**：边界框必须完全包含目标元素的所有可见部分，宁可稍大也不能遗漏任何部分
- **紧密贴合**：边界框应尽可能紧密地贴合元素的实际轮廓，避免过度扩大（如覆盖超过50%图像区域的情况）[[5]]
- **避免偏移**：特别注意边界框的位置准确性，确保左上角(x_min, y_min)和右下角(x_max, y_max)坐标精确对应元素的实际位置，防止bounding box drift（边界框漂移）问题

## 坐标系统说明
- **坐标系**：x轴为水平方向，从左到右递增；y轴为垂直方向，从上到下递增
- **像素精度**：坐标值必须基于实际像素位置，确保定位的精确性
- **边界定义**：
  - x_min：元素最左侧可见像素的x坐标
  - y_min：元素最顶部可见像素的y坐标  
  - x_max：元素最右侧可见像素的x坐标
  - y_max：元素最底部可见像素的y坐标

## 元素选择标准
- 只对总结中明确提及的、具有视觉显著性的重要元素进行定位
- 优先选择：人物、关键物体、文字标识、重要动作发生的位置
- 避免对模糊、遮挡严重或不重要的背景元素进行定位

## 时间点准确性
- 记录元素**首次清晰可见**的时间点
- 时间点必须在视频切片的时间范围内（{current_frame_slice['start_seconds']}秒到{current_frame_slice['end_seconds']}秒）
- 确保该时间点的画面确实包含该元素，避免时间点错误

# 输出格式要求
最终输出必须严格遵循以下JSON Schema，确保结构化输出的准确性 

{{
  "understanding_result": "详细的视频内容总结，包含关键视觉元素的描述",
  "start_second": {current_frame_slice['start_seconds']},
  "end_second": {current_frame_slice['end_seconds']},
  "mark_items": [
    {{
      "name": "元素的准确名称和描述（如：'穿红色衣服的主持人'，'屏幕左上角的公司logo'）",
      "start_second": "该元素首次清晰出现的时间点",
      "bbox": [x_min, y_min, x_max, y_max]  // 整数数组，严格遵循坐标定义
    }}
  ]
}}

# 质量检查提醒
在输出前，请严格检查：
1. 每个边界框是否完全覆盖了对应元素
2. 坐标值是否在合理范围内（0 ≤ x_min < x_max ≤ 图像宽度，0 ≤ y_min < y_max ≤ 图像高度）
3. 时间点是否准确对应元素出现的时刻
4. 只包含真正重要且在总结中提及的元素
        """
        if prev_context:
            prompt += f"""
                之前的视频切片的分析结果为为：\n{prev_context},请作为上下文数据输入，请结合之前的视频切片的分析结果和当前视频的分析结果，整合出一份合并后的数据。
            """

        return prompt
